import sys
import os
from optparse import OptionParser
from lib import table

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
usagetxt = "usage: %prog  -d project_dir "
parser = OptionParser(usage=usagetxt, version="%prog 1.0")
parser.add_option("-d", "--dir", dest="dir",
                  help="folder with the seqCluster project folder", metavar="DIR")


def print_iduinfo(info):
	for idu in info.keys():
		print idu
		for s in info[idu].sample.keys():
			print "\t%s" % s
			for i in  info[idu].getids(s):
				print "\t%s" % i
				for l in sampleinfo[s].getloci(i):
					print "\t\t%s" % l


def print_iduinfo_exp_table(info):
	out="idu"
	sl=sorted(sampleinfo.keys())
	for s in sl:
		out = out + "\t%s" % s
	out=out+"\n" 
	for idu in info.keys():
		out=out+str(idu)
		for s in sl:
			#print "\t%s" % s
			exp=0
			if info[idu].sample.has_key(s):

				for i in  info[idu].getids(s):
					#print "\t%s" % i
					exp+=sampleinfo[s].id[i]
			out=out+"\t%s" % str(exp)
		out=out+"\n"
	return out

def print_iduinfo_exp_html(info,tableidu):
	out=table.make_cell_header("idu")
	sl=sorted(sampleinfo.keys())
	for s in sl:
		out = out + table.make_cell_header(s)
	content=table.make_header(out)
	for idu in info.keys():
		out=""
		idname="C:"+str(idu)
		out=table.make_cell(table.make_a(table.make_link(idname,"ids.html#"+idname),idname))
		for s in sl:
			#print "\t%s" % s
			exp=0
			if info[idu].sample.has_key(s):
				
				for i in  info[idu].getids(s):
					#print "\t%s" % i
					exp+=sampleinfo[s].id[i]
			out=out+table.make_cell(str(exp))
		content=content+table.make_line(out)
	content=table.make_jshtml(table.make_table(content,"tableexp"),"tableexp")

	return content


def print_iduinfo_table(info):
	out="idu"
	sl=sorted(sampleinfo.keys())
	for s in sl:
		out = out + "\t%s" % s
	out=out+"\n" 
	for idu in info.keys():
		out=out+"C:"+str(idu)
		for s in sl:
			#print "\t%s" % s
			exp=0
			listid=""
			if info[idu].sample.has_key(s):
				
				for i in  info[idu].getids(s):
					#print "\t%s" % i
					listid=listid+";"+str(i)
			out=out+"\t%s" % listid
		out=out+"\n"
	return out


def print_iduinfo_table_html(info,tableexp):
	out=table.make_cell_header("idu")
	sl=sorted(sampleinfo.keys())
	for s in sl:
		out = out + table.make_cell_header(s)
	content=table.make_header(out)
	for idu in info.keys():
		out=""
		idname="C:"+str(idu)
		out=table.make_cell(table.make_a(table.make_link(idname,"exp.html#"+idname),idname))
		for s in sl:
			#print "\t%s" % s
			tablesample="../%s/html/info.html" % (s)
			exp=0
			listid=""
			if info[idu].sample.has_key(s):
				
				for i in  info[idu].getids(s):
					#print "\t%s" % i
					listid=listid+";"+table.make_link(str(i),tablesample+"#C:"+str(i))
			out=out+table.make_cell(listid)
		content=content+table.make_line(out)
	content=table.make_jshtml(table.make_table(content,"tableid"),"tableid")

	return content


def parse_merge(file,idv):
	f = open(file, 'r')
	newpos=[]
	for line in f:
		line=line.strip()
		c = line.split("\t")
		newpos.append("%s\t%s\t%s\t%s\t%s\t%s\n" % (c[0],c[1],c[2],idv,0,c[3]))
		#print line
	return newpos

def update_idu(new,old,s2):
	now={}
	#get all idu from new
	#print_iduinfo(new)
	for newidu in new.keys():
		now[newidu]=samplelist()
		#for each new-idu, get all pre-idus
		if new[newidu].sample.has_key('merge'):
			for preidu in new[newidu].getids('merge'):
				#get all ids in old with pre-idus
				for s in old[preidu].sample.keys():
					now[newidu].addsample(s)
					for i in old[preidu].getids(s):
						now[newidu].addids(s,i)
		if new[newidu].sample.has_key(s2):
			now[newidu].addsample(s2)
			for i in new[newidu].getids(s2):
				now[newidu].addids(s2,i)
	return now

		#put all these ids in now[new-idu]
class clusterclass:
    def __init__(self):
        self.id={}
        self.idl={}

    def test(self):
    	return 0

    def addidexp(self,ids,exp):
		self.id[ids]=exp
		return 0

    def addidl(self,ids,l):
		if (self.idl.has_key(ids)):
			llist=self.idl[ids]
			llist[l]=1
			self.idl[ids]=llist
		else:
			llist={}
			llist[l]=1
			self.idl[ids]=llist
		return 0

    def isid(self,ids):
		if (self.idl[ids]):
			return 0
    def getloci(self,ids):
		return self.idl[ids]

class idlist:
	def __init__(self):
		self.ids={}
	def addids(self,ids):
		self.ids[ids]=0
	def getids(self):
		return self.ids.keys()


class samplelist:
	def __init__(self):
		self.sample={}
	def addsample(self,s):
		#print "add sample %s " % s 
		self.sample[s]=idlist()
	def addids(self,s,ids):
		self.sample[s].addids(ids)
	def getids(self,s):
		return self.sample[s].getids()



class common_class:
	def __init__(self,idu):
		self.sample={}
		self.sample2={}
		self.common={}
		self.idv=idu
	def addsample(self,s):
		id2idu={}
		self.sample[s]=id2idu
	def addsampleid(self,s, ids,idu):
	 	if (self.sample.has_key(s)):
			self.idv=idu
			id2idu=self.sample[s]
			id2idu[ids]=idu
			self.sample[s]=id2idu
			
	def newiducommon(self,idu):
		#print "new idu %s " % idu
		self.common[idu]=samplelist()
	def addsample2idu(self,idu,s):
		self.common[idu].addsample(s)
	def addiducommon(self,idu,s,ids):
		self.common[idu].addids(s,ids)
	def getallidu(self):
		return self.common.keys()
	def getidsfromidu(self,idu,s):
		return self.common[idu].getids(s)

	def isid(self,s,ids):
		if (self.sample.has_key(s)):
			idsample=self.sample[s]
			if (idsample.has_key(ids)):
				return 1
	def getidu(self,s,ids):
			id2idu=self.sample[s]
			return id2idu[ids]	
	def issample(self,s):
		if (self.sample.has_key(s)):
			return 1
		

def print_common(common_info):
	for i in common_info.common.keys():
		print i
		for s in common_info.common[i].sample.keys():
			print "\t%s" % s
			for ids in common_info.sample[s]:
				print "\t%s" % ids

def merge_cross(file1,idv,s1,s2):
	slist=common_class(0)
	slist.addsample(s1)
	slist.addsample(s2)
	
	f = open(file1, 'r')
	idv=idv
	clus={}
	for line in f:
		line=line.strip()
		cols = line.split("\t")
		c1id=int(cols[3])
		c2id=int(cols[9])
		exp1=cols[4]
		exp2=cols[10]
		#print slist.issample(s1) 
		if (slist.issample(s1) and slist.issample(s2)):
			#print "samples in list %s %s" % (s1,s2)
			#check that idu are always the same when exists and assign many time to the same cluster
			#if (c2id==37):
			#	print line
			if (slist.isid(s1,c1id)):
				idu=slist.getidu(s1,c1id)
				slist.addsampleid(s2,c2id,idu)
				slist.addiducommon(idu,s2,c2id)
				#print "new cluster %s  in s2  %s" % (idv,c2id)
			elif (slist.isid(s2,c2id)):
				idu=slist.getidu(s2,c2id)
				slist.addsampleid(s1,c1id,idu)
				slist.addiducommon(idu,s1,c1id)
				#print "new cluster %s  in s1 %s" % (idv,c1id) 
			else:
				idv+=1
				#print "new cluster %s  in both %s %s" % (idv,c1id,c2id) 
				
				slist.addsampleid(s1,c1id,idv)
				slist.addsampleid(s2,c2id,idv)
				slist.newiducommon(idv)
				slist.addsample2idu(idv,s1)
				slist.addsample2idu(idv,s2)
				slist.addiducommon(idv,s1,c1id)
				slist.addiducommon(idv,s2,c2id)
	slist.idv=idv
	return slist



if len(sys.argv)==1:
    parser.print_help()
    sys.exit(1)
    
(options, args) = parser.parse_args()

if (options.dir[-1]=="/"):
	options.dir=options.dir[:-1]

dirs = sorted(os.listdir(options.dir+"/"))

index=0
tempfile1=options.dir+"/t1.bed"
tempfile2=options.dir+"/t2.bed"
tempfile3=options.dir+"/out.bed"
tempfile4=options.dir+"/merge.bed"
output=options.dir+"/table_%s.txt"
sampleinfo={}

#iduinfo_old= None
#add to sample_class merge sample
for dn in dirs:
	#print d
	d=options.dir+"/"+dn
	if os.path.isdir(d):
		clusfile=d+"/clus.parse.txt"
		#print d 
		if (os.path.exists(clusfile)):
			index+=1
			iduinfo={}
			print clusfile+"\n"
			tempdirs=clusfile.split("/")
			names=tempdirs[-2]
			idv=0
			if (index==1):
				out = open(tempfile1, 'w')
				
				names1=names
				
				#add to sample_class both samples
			else:
				out = open(tempfile2, 'w')
				#add to sample_class both samples
			
			cluster=clusterclass()
			print "name sample %s\n" % names 
			f = open(clusfile, 'r')
			for line in f:
				line=line.strip()
				cols = line.split(" ")
				#print cols[0]
				if (cols[0]=="C"):
					exp=0
					clus=cols[1]
				if (cols[0]=="S"):
					exp=exp+int(cols[2])
				if (cols[0]=="L"):
					locu="%s\t%s\t%s\t%s\t%s\t%s\n" % (cols[2],cols[3],cols[4],clus,0,cols[5])
					out.write(locu)
					#print cluster.test()
					cluster.addidexp(int(clus),exp)
					cluster.addidl(int(clus),locu)
			
			sampleinfo[names]=cluster
			out.close()

			if (index>1):
				print index	

				#intersectBed files
				cmd2intersect="intersectBed -wo -a %s -b %s > %s" % (tempfile1,tempfile2,tempfile3)
				print bcolors.OKBLUE+cmd2intersect+bcolors.ENDC
				os.system(cmd2intersect)
				#sys.exit(1)
				clustert=clusterclass()
				if (index>2):
					names1="merge"
				common_info=merge_cross(tempfile3,idv,names1,names)
				#specific check

				#print_common(common_info)
				idv=common_info.idv
				#loci in sample1, loci in sample2, loci in merge
				#check s1
				out = open(tempfile1, 'w')
				#cluster first sample (s1 or merge)
				#print "adding sample 1"
				for idc in cluster1.id.keys():

					if not common_info.isid(names1,idc):
						#print "idc %s of merge is not in common " % idc
						loci=cluster1.getloci(idc)
						idv+=1
						#if iduinfo_old:
						iduinfo[idv]=samplelist()
						iduinfo[idv].addsample(names1)
						iduinfo[idv].addids(names1,idc)
						for l in loci.keys():
							#be carefull with diff lines with same cluster id
							#print l
							c=l.split("\t")
							c[3]=str(idv)
							l="\t".join(c[0:6])
							clustert.addidexp(idv,0)
							clustert.addidl(idv,l)
							#if idv==1213:
								#print "sample 1 %s" % l
							out.write(l+"\n")
				#common samples	
				#print "\n\noverlap loci in clusters 1 and 2"			
				for idu in common_info.getallidu():
					out2 = open(tempfile4, 'w')
					#get for names1 the ids
					allloci=[]
					iduinfo[idu]=samplelist()
					iduinfo[idu].addsample(names1)
					for ids in common_info.getidsfromidu(idu,names1):
						iduinfo[idu].addids(names1,ids)
						for l in cluster1.getloci(ids):
							c=l.split("\t")
							out2.write("%s\t%s\t%s\t%s\t%s\t%s\n" % (c[0],c[1],c[2],idu,0,c[5]))
							#if idu==1213:
								#print "common sample 2 %s" % l
					iduinfo[idu].addsample(names)
					for ids in common_info.getidsfromidu(idu,names):
						iduinfo[idu].addids(names,ids)
						for l in cluster.getloci(ids):
							c=l.split("\t")
							out2.write("%s\t%s\t%s\t%s\t%s\t%s\n" % (c[0],c[1],c[2],idu,0,c[5]))
							#if idu==1213:
								#print "common sample 2 %s" % l
					out2.close()
					cmd2intersect="mergeBed -s -i %s > %s" % (tempfile4,tempfile3)
					#print bcolors.OKBLUE+cmd2intersect+bcolors.ENDC
					os.system(cmd2intersect)

					#overlap common loci to get common regions that will be the loci for this idu
					newpositions=parse_merge(tempfile3,idu)
					for l in newpositions:
						clustert.addidexp(idu,0)
						clustert.addidl(idu,l)
						#print l
						out.write(l)
					#update idu<=>ids relations when index>2
				
				#cluster second samples
				#print "adding sample 2"
				for idc in cluster.id.keys():
					if not common_info.isid(names,idc):
						loci=cluster.getloci(idc)
						idv+=1
						iduinfo[idv]=samplelist()
						iduinfo[idv].addsample(names)
						iduinfo[idv].addids(names,idc)
						for l in loci.keys():
							#be carefull with diff lines with same cluster id
							c=l.split("\t")
							c[3]=str(idv)
							l="\t".join(c[0:6])
							clustert.addidexp(idv,0)
							clustert.addidl(idv,l)
							#if idv==1213:
								#print "sample 2 %s" % l
							out.write(l+"\n")
				out.close()

				if index>2:
					#print "doing update"
					#print "checking cluster 1213"
					#for s in iduinfo_old[1213].sample.keys():
					#	print s
					#	for i in iduinfo_old[1213].getids(s):
					#		print i

					iduinfo=update_idu(iduinfo,iduinfo_old,names)

				
				
				iduinfo_old=iduinfo
				cluster1=clustert
				#print "idu and sample and ids in old iduinfo"
				
			else:
				cluster1=cluster
				#call function to get the common and news clusters: if one overlap, is count as the same cluster

out = open(output % "exp", 'w')
#out.write(print_iduinfo(iduinfo_old))
out.write(print_iduinfo_exp_table(iduinfo_old))
out.close()
out = open(output % "ids", 'w')
#out.write(print_iduinfo(iduinfo_old))
out.write(print_iduinfo_table(iduinfo_old))
out.close()

#creating html
cmd="mkdir %s" % (options.dir+"/html")
print bcolors.OKBLUE+cmd+bcolors.ENDC 
os.system(cmd)

pathscript=os.path.abspath(__file__).replace("merge.samples.py","")
cmdcp="rsync -a -u --exclude='- *.' "+pathscript+"js "+pathscript+"css "+pathscript+"images "+options.dir+"/html/."
print bcolors.OKBLUE+cmdcp+bcolors.ENDC 
os.system(cmdcp)

out = open(options.dir+"/html/ids.html", 'w')
out.write(print_iduinfo_table_html(iduinfo_old,output % "exp.html"))
out.close()

out = open(options.dir+"/html/exp.html", 'w')
out.write(print_iduinfo_exp_html(iduinfo_old,output % "ids.html"))
out.close()
